package com.shuke.my_test.shardingSphere.manager;


import com.shuke.my_test.shardingSphere.domain.UserInfo;
import com.shuke.my_test.shardingSphere.service.UserInfoService;
import com.zaxxer.hikari.pool.HikariProxyConnection;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.rule.ShardingRule;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @author 舒克、舒克
 * @date 2025/7/25 10:48
 * @description 动态分片管理类 修改availableTargetNames
 */
@Component
@Slf4j
public class DynamicShardingManager {

    private static final String LOGIC_TABLE_NAME = "user_info";

    private static final String LOGIC_DB_NAME = "my_test";

    @Resource
    private UserInfoService userInfoService;

    @Resource
    private DataSource dataSource;

    @PostConstruct
    public void initialize(){
        log.info("初始化动态分表配置");
        updateShardingTable();
    }

    /**
     * 更新分表配置
     */
    private void updateShardingTable() {
        Set <String> tableNames = getAllTableNames();
        String newTableNames = tableNames.stream()
                .map(tableName -> LOGIC_DB_NAME + "." + tableName)
                .collect(Collectors.joining(","));
        log.info("动态分表的表名：{}",newTableNames);

        // 获取 ShardingSphere  的 上下文对象
        ContextManager shardingSphereContext =getShardingSphereContext();
        // 从上下文对象中获取 规则元数据
        ShardingSphereRuleMetaData ruleMetaData = shardingSphereContext.getMetaDataContexts()
                .getMetaData()
                .getDatabase(LOGIC_DB_NAME)
                .getRuleMetaData();

        Optional<ShardingRule> singleRule = ruleMetaData.findSingleRule(ShardingRule.class);

        if (singleRule.isPresent()){
            ShardingRuleConfiguration configuration = (ShardingRuleConfiguration)singleRule.get().getConfiguration();

            List<ShardingTableRuleConfiguration> rules = configuration.getTables()
                    .stream()
                    .map(rule -> {
                        if(LOGIC_TABLE_NAME.equals(rule.getLogicTable())){
                            ShardingTableRuleConfiguration newRule = new ShardingTableRuleConfiguration(LOGIC_TABLE_NAME, newTableNames);
                            newRule.setDatabaseShardingStrategy(rule.getDatabaseShardingStrategy());
                            newRule.setTableShardingStrategy(rule.getTableShardingStrategy());
                            newRule.setKeyGenerateStrategy(rule.getKeyGenerateStrategy());
                            newRule.setAuditStrategy(rule.getAuditStrategy());
                            return newRule;
                        }
                        return rule;
                    })
                    .toList();
            configuration.setTables(rules);
            shardingSphereContext.alterRuleConfiguration(LOGIC_DB_NAME, List.of(configuration));
            shardingSphereContext.reloadDatabase(LOGIC_DB_NAME);
            log.info("动态分表的表名更新成功");
        }else{
            log.error("未找到分片规则");
        }

    }

    /**
     * 获取 ShardingSphere  的 上下文对象
     * @return
     */
    private ContextManager getShardingSphereContext() {
        try {
            ShardingSphereConnection connection = dataSource.getConnection().unwrap(ShardingSphereConnection.class);
            return connection.getContextManager();
        } catch (SQLException e) {
            throw new RuntimeException("获取 ShardingSphere 的 上下文对象失败：",e);
        }
    }

    /**
     *  获取所有的表名 包括初始表 和 分表
     * @return
     */
    private Set<String> getAllTableNames() {
        Set<Long> ids = userInfoService.list().stream().map(UserInfo::getId).collect(Collectors.toSet());
        Set<String> tableNames = ids.stream().map(id -> LOGIC_TABLE_NAME + "_" + id).collect(Collectors.toSet());
        tableNames.add(LOGIC_TABLE_NAME);
        return tableNames;
    }
}
